# pre_download.py - 预下载最小模型
from huggingface_hub import snapshot_download
import os

# 设置使用国内镜像加速下载
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 选择最小模型
model_id = "sentence-transformers/all-MiniLM-L6-v2"
local_dir = "./models/all-MiniLM-L6-v2"

print(f"开始下载模型: {model_id}")
print("模型大小: 约 80MB")

try:
    # 下载模型到本地目录
    model_path = snapshot_download(
        repo_id=model_id,
        local_dir=local_dir,
        local_dir_use_symlinks=False
    )
    print(f"模型已成功下载到: {model_path}")
except Exception as e:
    print(f"下载失败: {e}")
    print("请检查网络连接或尝试使用代理")